#!/usr/bin/env python2

# Network Security Guard (botnet)
#
# Copyright (C) 2010-2011 NSG Project
# Author: Choonho Son <choonho@kt.com>
# URL: <http://nsgbot.googlecode.com>
# For license information, see LICENSE.TXT
#

import pcap
import dpkt
import dnet
import sys
import string
import socket
import re
import time
import Queue
import threading
import signal
import struct

from threading import Thread
from optparse import OptionParser


# global variable
queue_count = 0                        # total analyzed packet count
ExIP = []                              # Exception IPs for capturing
DADDR = socket.inet_aton('127.0.0.1')  # Capturing IP
TCP_SIG = {}                           # TCP signature table
UDP_SIG = {}                           # UDP signature table
dev = ""
verbose = False

# make each Queue for thread
queue_packet = Queue.Queue()
queue_tcp = Queue.Queue()
queue_udp = Queue.Queue()
queue_etc = Queue.Queue()

# SIGNAL
WORKING = True
LOCK = threading.Lock()
STOP = 0

class Signalled(Exception):
    pass

def sigBreak(signum, f):
    global STOP
    LOCK.acquire()
    STOP = 1
    LOCK.release()
    raise Signalled

def softBreak():
    global STOP
    LOCK.acquire()
    STOP = 1
    LOCK.release()
    raise Signalled

def print_packet(pktlen, data, timestamp):
    if not data:
        return

    eth = dpkt.ethernet.Ethernet(data)

    if eth.type == dpkt.ethernet.ETH_TYPE_IP:
        ip = eth.data
    else:
        return

    # Check exception IPs
    for exip in ExIP:
        if ip.src == exip or ip.dst == exip:
            return

    # Dst IP must be DADDR
    # since, there are a lot of multicast traffic
    global queue_count
    if ip.dst == DADDR and ip.p == dpkt.ip.IP_PROTO_TCP:
        # Push to wanted Queue
        queue_count += 1
        queue_tcp.put(ip)
    elif ip.dst == DADDR and ip.p == dpkt.ip.IP_PROTO_UDP:
        # Push to UDP Queue
        queue_udp.put(ip)
        queue_count += 1

    elif ip.dst == DADDR:
        # exclude TCP/UDP 
        print "[ETC] %s(%s)" % (socket.inet_ntoa(ip.src), ip.p)
        #queue_etc.put(ip)
        #queue_count += 1
    return



class Packet_Parser(Thread):
    def run(self):
        print "Start Packet Parser.."
        p = pcap.pcapObject()
        #dev = sys.argv[1]
        print "dev:",dev
        p.open_live(dev, 1600, 1, 100)
        
        while STOP == 0:
            LOCK.acquire()
            stop = STOP
            LOCK.release()
            if stop:
                print '%s' % sys.exc_type
                print 'shutting down'
                print '%d packets receive, %d packets dropped, %d packets dropped by interface' % p.stats()

                # send False to each queue for stop
                queue_tcp.put(False)

                break

            p.dispatch(1, print_packet)                
        print "Finished Packet_Parser"

class ParseTCP(Thread):
    def __init__(self, queue):
        threading.Thread.__init__(self)
        self.queue = queue
        self.sock = dnet.ip()

    def run(self):
        print "Start TCP Parse...."
        # get IP/TCP packet
        decoded = self.queue.get()
        while STOP == 0:
            self.parse_tcp(decoded)
            decoded = self.queue.get()
        debug("Finished ParseTCP")

    def parse_tcp(self, ip):
        tcp = ip.data
        # filter ssh port
        if tcp.dport == 22:
            return

        print_tcp(ip, queue_count) 
        key = (ip.src, tcp.sport, ip.dst, tcp.dport)
        if TCP_SIG.has_key(key) == False:
            # check TCP SYN flags
            if (tcp.flags & dpkt.tcp.TH_SYN) > 0: # TCP SYN (1 segment)
                result = {'orgseq':tcp.seq, 'client.seq': tcp.seq, 'server.seq': 0, \
                              'signature':{tcp.seq:tcp.data},'ftime':time.time()}
                TCP_SIG[key] = result
                print_tcp(ip, "Start TCP Connection", debug=1)
            else:
                #print "Not first TCP packet"
                return

        # key always exist
        result = TCP_SIG[key]

        if (tcp.flags & dpkt.tcp.TH_SYN) > 0: # TCP SYN
            # reply to TCP SYN
            self.send_SYNACK(ip)
            result['client.seq'] = tcp.ack 
            result['server.seq'] = 1
            print_tcp(ip, "2 Segment") 
        else:
            # reply to TCP ACK
            received_seq  = tcp.seq
            wanted_seq = result['client.seq']
            result['server.seq'] = tcp.ack
            result['ltime'] = time.time()
            #print "recv seq:%d vs wanted seq:%d" % (received_seq, wanted_seq)
            if received_seq == wanted_seq:
                sig = result['signature']
                sig[received_seq] = tcp.data
                result['client.seq'] = result['client.seq'] + len(tcp.data) 
                #print "Signature:",sig
                #print TCP_SIG
                #show_TCP_SIG()
                # get NEXT SEQ
                seq = result['client.seq']
                self.send_ACK(ip, seq, tcp.ack)
                print_tcp(ip, "Data Segment(continue)")
            elif received_seq > wanted_seq:
                self.send_ACK(ip, wanted_seq,tcp.ack)
                print_tcp(ip, "Data Segment(retrans)")
            #else:
                #self.send_ACK(ip, wanted_seq)
                #print_tcp(ip, "Data Segment(next)")
        #print "Analyzed packet count: ", queue_count
        #print_tcp(ip)
                
    def send_SYNACK(self, ip):
        tcp = ip.data
        tcp.flags = dpkt.tcp.TH_SYN | dpkt.tcp.TH_ACK
        ip.src, ip.dst = ip.dst, ip.src
        tcp.seq, tcp.ack = 0, tcp.seq+1
        tcp.sport, tcp.dport = tcp.dport, tcp.sport
        self.send_tcp(ip, tcp)

    def send_ACK(self, ip, seq, ack=1):
        tcp = ip.data
        tcp.flags = tcp.flags | dpkt.tcp.TH_ACK
        ip.src, ip.dst = ip.dst, ip.src
        tcp.seq, tcp.ack = ack, seq 
        tcp.sport, tcp.dport = tcp.dport, tcp.sport
        self.send_tcp(ip, tcp)

    def send_tcp(self, ip, tcp, ttl=50):
        #tcp.data = struct.pack('H',0)
        self.sock = dnet.ip()
        ip.ttl = ttl
        buf = dnet.ip_checksum(str(ip))
        self.sock.send(buf)

    def parse_udp(self, ip):
        pass


    def print_ip(self, ip):
        layer4 = ip.data
        print "[+]GET %s(%d)-(%d)->%s(%d):(flag:%d,SEQ:%d,ACK:%d)%s" % (socket.inet_ntoa(ip.src), layer4.sport, ip.p, socket.inet_ntoa(ip.dst), layer4.dport, layer4.flags, layer4.seq, layer4.ack, layer4.data)

#######################
# UDP signature
#######################
class ParseUDP(Thread):
    def __init__(self, queue):
        threading.Thread.__init__(self)
        self.queue = queue

    def run(self):
        print "Start UDP Parse...."
        # get IP/UDP packet
        decoded = self.queue.get()
        while STOP==0:
            self.parse_udp(decoded)
            decoded = self.queue.get()
        debug("Finished ParseUDP")

    def parse_udp(self, ip):
        udp = ip.data
        key = (ip.src, udp.sport, ip.dst, udp.dport)
        if UDP_SIG.has_key(key) == False:
            # First UDP data
            result = {'ftime':time.time(), 'signature': [udp.data] }
            UDP_SIG[key] = result
        else:
            # Append UDP data
            result = UDP_SIG[key]
            result['ftime'] = time.time()
            data = result['signature'] 
            data.append(udp.data)
            
        print_udp(ip, debug=1)


def print_tcp(ip,str="+", debug=0):
    if debug == 0:
        return
    
    layer4 = ip.data
    print "[TCP %s] %s(%d)-(%d)->%s(%d):(flag:%d,SEQ:%d,ACK:%d)%s" % \
        (str,socket.inet_ntoa(ip.src), layer4.sport, ip.p, socket.inet_ntoa(ip.dst), layer4.dport, \
             layer4.flags, layer4.seq, layer4.ack, layer4.data)
    
def print_udp(ip,str="+", debug=0):
    if debug == 0:
        return
    udp = ip.data
    print "[UDP %s]GET %s(%d)-(%d)->%s(%d):%s" % (str, socket.inet_ntoa(ip.src), udp.sport, ip.p, socket.inet_ntoa(ip.dst), udp.dport, udp.data)

def startAnalyzer():

    parser = OptionParser()
    parser.add_option("-i", "--interface", dest="interface", help="Monitoring Network Interface", default="eth0")
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose", help="Debug options")
    parser.add_option("-d", "--daddr", dest="daddr", help="Destination IP", metavar="127.0.0.1")
    parser.add_option("-e", "--exip", dest="exip", help="Exception IP")

    (options, args) = parser.parse_args()

    # configure global variables
    global dev
    global DADDR
    global ExIP
    global verbose

    if options.verbose:
        verbose = True

    # Network Interface
    if options.interface:
        dev = options.interface
    else:
        dev = "eth0"
 
    # Destination IP
    if options.daddr:
        debug("Destination IP: %s" % options.daddr)
        DADDR = socket.inet_aton(options.daddr)
    # Filter IPs (exceptions)
    if options.exip:
        debug("Append Exception IP: %s" % options.exip)
        ExIP.append(socket.inet_aton(options.exip))

    #print blackList
    thr_par = Packet_Parser()
    thr_tcp = ParseTCP(queue_tcp)
    thr_udp = ParseUDP(queue_udp)
    thr_con = Console()
    
    thr_par.start()
    thr_tcp.start()
    thr_udp.start()
    thr_con.start()
    try:
        signal.pause()
    except Signalled:
        print "* Singalled"

    print "queue join"

    print "finish queue join"
    thr_par.join()
    thr_tcp.join()
    thr_udp.join()
    thr_con.join()

class Console(Thread):
    def run(self):
        print "Start Console mode:"
        while STOP == 0:
            cmd = raw_input("Console(help:?) >")
            if cmd == "Q" or cmd=="exit" or cmd == "q":
                softBreak()
                return

            index = cmd.split(" ")
            if len(index) != 3:
                self.showHelp()
                continue
            # Cmd parser
            if index[0] == "tcp":
                show_TCP_SIG(index[1], index[2])

            elif index[0] == "?":
                self.showHelp()
        debug("Finished Console")

    def showHelp(self):
        msg = """

usage: <protocol> <depth> <filter>
   ex: tcp 0 123.123.123.1

<protocol> : [tcp|udp]
<depth>    : [0|1]  , 0 shows only ip and port, 1 shows signature
<filter>   : ip, which want to see , 'all' means show ALL

"""
        print msg

def show_TCP_SIG(depth, filter):
    count = len(TCP_SIG.keys())
    print "="*20 , "TCP Signature(%d)" % count, "="*20
    output = ""
    for key in TCP_SIG.keys():
        result = TCP_SIG[key]
        (saddr, sport, daddr, dport) = key
        if filter == "all":
            output = output + tcp_output(key, result, depth)
            # need filter
        elif socket.inet_ntoa(saddr) == filter:
            output = output + tcp_output(key, result, depth)

    print output


def tcp_output(key, result, depth):
    (saddr, sport, daddr, dport) = key
    if depth == "0":
        return "%s(%s)->%s(%s)\n" % (socket.inet_ntoa(saddr), sport, socket.inet_ntoa(daddr), dport)        
    elif depth == "1":
        return "%s(%s)->%s(%s) : %s\n" % (socket.inet_ntoa(saddr), sport, socket.inet_ntoa(daddr), dport, toSignature(result['signature']))
    else:
        return ""
    
def toSignature(dic):
    keys = dic.keys()
    keys.sort()
    result = ""
    for key in keys:
        result = result + dic[key]
    return result

def debug(str):
    if verbose == True:
        print "[DEBUG] ", str

if __name__ == "__main__":

    signal.signal(signal.SIGINT, sigBreak)
    # packet capture & analyze
    startAnalyzer()
